# deepseek_v3_demo.py

import os
import subprocess
import time
from pathlib import Path

import modal
from modal.forward import forward

# --- 1. 配置 ---
# Hugging Face 模型仓库 ID
MODEL_ID = "deepseek-ai/DeepSeek-V3.1"
# 用于缓存模型的本地目录路径
MODEL_DIR = "/model_cache"
# GPU 配置：至少需要 8 颗 H100 (40GB) GPU。如果模型需要更高精度，则需 80GB 版本。
# 请注意：这是一个非常昂贵的配置！
GPU_CONFIG = modal.gpu.H100(count=8)

# --- 2. 定义容器镜像 ---
# 我们将构建一个包含所有必要依赖的 Docker 镜像。
vllm_image = (
    modal.Image.from_registry("nvidia/cuda:12.1.1-devel-ubuntu22.04", add_python="3.11")
    .pip_install(
        "vllm==0.4.2",
        "huggingface_hub",
        "hf-transfer",
        "torch==2.1.2",
        "transformers==4.41.2",
        "fastapi",
        "uvicorn",
        "httpx",
    )
    # 使用更快的 hf-transfer 下载模型
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

# --- 3. 创建 Modal App 和持久化存储 ---
app = modal.App("deepseek-v3-1-server", image=vllm_image)
# 创建一个持久化的 Volume 来缓存模型权重
model_volume = modal.Volume.from_name(
    "deepseek-v3-model-cache", create_if_missing=True
)

# --- 4. 定义 vLLM 推理服务器 ---
@app.cls(
    gpu=GPU_CONFIG,
    volumes={MODEL_DIR: model_volume},
    # 容器闲置 5 分钟后自动关闭以节约成本
    container_idle_timeout=300,
    # 允许处理大量并发请求
    allow_concurrent_inputs=100,
    # 需要 Hugging Face 访问令牌
    secrets=[modal.Secret.from_name("huggingface-secret")],
)
class VLLMServer:
    """
    一个管理 vLLM 服务器生命周期的有状态类。
    """
    @modal.enter()
    def start_vllm_server(self):
        """
        容器启动时执行，下载模型并启动 vLLM 服务器。
        """
        vllm_port = 8000

        # vLLM 启动命令
        # --tensor-parallel-size 8: 这是关键，表示使用 8 颗 GPU 进行张量并行
        # --max-model-len 4096: 可根据需求调整上下文长度
        cmd = [
            "python", "-m", "vllm.entrypoints.openai.api_server",
            "--host", "0.0.0.0",
            "--port", str(vllm_port),
            "--model", MODEL_ID,
            "--tensor-parallel-size", str(GPU_CONFIG.count),
            "--download-dir", MODEL_DIR,
            "--max-model-len", "4096",
            "--dtype", "auto",
            # 如果遇到内存不足问题，可以尝试 gpu_memory_utilization
            # "--gpu-memory-utilization", "0.90",
        ]

        print("🚀 正在启动 vLLM 服务器...")
        # 在后台启动 vLLM 进程
        self.vllm_process = subprocess.Popen(cmd)

        # 等待 vLLM 服务器准备就绪
        # 通过轮询其健康检查端点来确认
        start_time = time.time()
        while True:
            if time.time() - start_time > 15 * 60: # 15 分钟超时
                raise RuntimeError("vLLM 服务器启动超时！")
            
            try:
                # 使用 httpx 检查健康状态
                import httpx
                response = httpx.get(f"http://localhost:{vllm_port}/health")
                if response.status_code == 200:
                    print("✅ vLLM 服务器已准备就绪！")
                    break
            except (httpx.ConnectError, httpx.ReadError):
                # 服务器尚未启动，等待一会
                print("⏳ 等待 vLLM 服务器响应...")
                time.sleep(10)
        
        # 将 vLLM 端口转发到本地，方便调试 (可选)
        # 运行 `modal serve` 时，你可以通过 localhost:8000 访问 vLLM
        # forward(vllm_port, self.vllm_process.pid)

    @modal.exit()
    def stop_vllm_server(self):
        """容器退出时，优雅地关闭 vLLM 进程。"""
        print("🛑 正在关闭 vLLM 服务器...")
        self.vllm_process.terminate()
        self.vllm_process.wait()
        print("✅ vLLM 服务器已关闭。")

    @modal.asgi_app()
    def api(self):
        """
        创建一个 FastAPI 应用作为代理，将所有请求转发到在同一容器中运行的 vLLM 服务器。
        这是实现流式响应和优雅处理请求的最佳实践。
        """
        from fastapi import FastAPI, Request, Response
        import httpx

        app = FastAPI()
        client = httpx.AsyncClient()

        @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
        async def proxy(request: Request, path: str):
            # 将客户端的请求原封不动地转发到 vLLM 服务器
            url = httpx.URL(f"http://localhost:8000/{path}")
            headers = dict(request.headers)
            # Host 头需要更新
            headers["host"] = "localhost"
            
            # 使用流式请求和响应
            req = client.build_request(
                request.method,
                url,
                headers=headers,
                params=request.query_params,
                content=request.stream(),
            )
            
            r = await client.send(req, stream=True)
            
            return Response(
                content=r.aiter_raw(),
                status_code=r.status_code,
                headers=dict(r.headers),
            )
        return app


# --- 5. 本地客户端，用于测试 API ---
@app.local_entrypoint()
def main():
    """
    一个本地入口点，用于与部署在 Modal 上的 API 进行交互。
    """
    import openai

    # 获取部署的 Web 服务的 URL
    # VLLMServer.api.web_url 在 `modal run` 或 `modal serve` 期间不可用，
    # 需要先 `modal deploy deepseek_v3_demo.py`
    try:
        url = VLLMServer.api.web_url
        print(f"连接到已部署的服务器: {url}")
    except Exception:
        print("请先使用 'modal deploy deepseek_v3_demo.py' 部署应用。")
        print("部署后，请从 Modal UI 获取 API 地址并替换下面的 'base_url'。")
        return

    client = openai.OpenAI(
        base_url=f"{url}/v1",  # vLLM 的 OpenAI 兼容 API 路径
        api_key="not-needed",  # vLLM 不需要 API 密钥
    )

    print("发送聊天请求...")

    messages = [
        {"role": "system", "content": "You are a helpful and creative assistant."},
        {"role": "user", "content": "Tell me a short story about a robot who discovers music."},
    ]

    try:
        # 使用流式响应
        stream = client.chat.completions.create(
            model=MODEL_ID,
            messages=messages,
            stream=True,
        )

        print("--- 机器人与旋律 ---")
        for chunk in stream:
            content = chunk.choices[0].delta.content
            if content:
                print(content, end="", flush=True)
        
        print("\n--- 流式响应结束 ---")

    except openai.APIConnectionError as e:
        print(f"无法连接到服务器: {e.__cause__}")
    except openai.APIStatusError as e:
        print(f"API 错误: 状态码 {e.status_code}, 响应: {e.response}")
    except Exception as e:
        print(f"发生未知错误: {e}")
